*~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
*	Compute OLS-VAM estimates
*	----------------------------------------------------------------------------
*	Syntax notes: 	varlist takes the form [Y] [enrollment dummies]
*				  	baseline controls go in [bcovs]
*					the var option also stores variance estimates
*					the if option estimates in a subset of the data but the
*						resulting estimates will still populate all obs
*	----------------------------------------------------------------------------

	program define olsvam, rclass
	program drop olsvam
	syntax varlist [if] [in], gen(name) [vce(string)] [covs(varlist fv)] [pscores(varlist fv)] [sample(varname)]

*~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~


*	preliminaries

	marksample touse

	gettoken Y D: varlist
	local X `covs'
	local P `pscores'
	local S `sample'

	//ivreg2 with partial option doesn't store omitted regressors so doing a check here
	_rmcoll `D' if `touse', noconstant forcedrop
	local newD r(varlist)
	local dropped: list D - newD
	return local dropped = "`dropped'"
	local D = r(varlist)

	unab Dvars: `D'
	local J: word count `Dvars'

	cap confirm var `name'
	if _rc == 0{
		di as error "Specified variable already exists in the dataset."
		exit 101
	}

	if !inlist("`vce'","","hom","het","hyb"){
		di as error `"VCE must be one of "hom", "het" or "hyb"'
		exit 101
	}

	if ("`sample'"!="") + ("`pscores'"!="") == 1{
		di as error "Must specify both sample and pscores, or neither."
		exit 101
	}

*	estimate
	
	tempname b b1 b2 b_short V C V_hom V_het T V_short R R1 R2

	//conventional
	if "`pscores'"==""{

		if "`vce'"!="het"{
			qui ivreg2 `Y' `D' `X' if `touse', partial(`X') nocons noid
			mat `V_hom' = e(V)
		}
		if inlist("`vce'","het","hyb"){
			qui ivreg2 `Y' `D' `X' if `touse', partial(`X') nocons noid r
			mat `V_het' = e(V)
		}

		if "`vce'"=="hom" mat `V' = `V_hom'
		if "`vce'"=="het" mat `V' = `V_het'
		if "`vce'"=="hyb"{
			local M = rowsof(`V_het')

			mat `V' = J(`M',`M',.)
			mat `C' = J(`M',`M',.)

			//correlation matrix (using homoscedastic corr. structure)
			forval i=1/`M'{
				forval j=1/`M'{
					qui mat `C'[`i',`j']=`V_hom'[`i',`j'] / sqrt(`V_hom'[`i',`i']*`V_hom'[`j',`j'])
				}
			}

			//fill diagonal
			forval i=1/`M'{
				qui mat `V'[`i',`i']=max(`V_hom'[`i',`i'],`V_het'[`i',`i'])
			}

			//fill off-diagonal
			forval i=1/`M'{
				forval j=1/`M'{
					if `i'!=`j' qui mat `V'[`i',`j']=`C'[`i',`j']*sqrt(`V'[`i',`i']*`V'[`j',`j'])
				}
			}
		}

		mat `b' = e(b)
	}


	//pscore controlled
	if "`pscores'"!=""{
		preserve

			qui keep if `touse'

			tempvar id
			gen `id' = _n
			//stack
			qui expand 2, gen(copy)
			//second copy is the lottery sample
			qui drop if copy == 1 & !`sample'
			//interact
			foreach var of varlist `D'{
				tempvar p_`var'
				qui gen `p_`var'' = `var' * (copy == 1)
				qui replace `var' = `var' * (copy == 0)
				local p_D `p_D' `p_`var''
			}
			//remove collinear
			qui _rmcoll `p_D' if copy == 1, noconstant forcedrop
			local p_D = r(varlist)
			qui _rmcoll `D' if copy == 1, noconstant forcedrop
			local lottD = r(varlist)

			//no pscores in OLS copy
			local Praw = subinstr("i.","`P'","",.)
			foreach var of varlist `Praw'{
				replace `var' = `var' * (copy == 1)
			}

			if inlist("`vce'","hom","hyb"){
				qui ivreg2 `Y' `D' `p_D' c.(`X' `P')#i.copy, partial(c.(`X' `P')#i.copy) nocons noid
				mat `b' = e(b)
				mat `V_hom' = e(V)
				local bnames `: colfullnames e(b)'
			}
			if inlist("`vce'","het","hyb") {
				qui ivreg2 `Y' `D' `p_D' c.(`X' `P')#i.copy, partial(c.(`X' `P')#i.copy) nocons noid cluster(`id')
				mat `b' = e(b)
				mat `V_het' = e(V)
				local bnames `: colfullnames e(b)'
			}
			if "`vce'"==""{
				qui ivreg2 `Y' `D' `X' if copy == 0, partial(`X') nocons noid
				mat `b1' = e(b)
				local b1names `: colfullnames e(b)'
				qui ivreg2 `Y' `p_D' `X' `P' if copy == 1, partial(`X' `P') nocons noid
				mat `b2' = e(b)
				local b2names `: colfullnames e(b)'
				mat `b' = [`b1',`b2']
				local bnames `b1names' `b2names'
			}

			if "`vce'"=="hom" mat `V' = `V_hom'
			if "`vce'"=="het" mat `V' = `V_het'
			if "`vce'"=="hyb"{
				local M = rowsof(`V_het')
				mat `V' = J(`M',`M',.)
				mat `C' = J(`M',`M',.)

				//correlation matrix (using homoscedastic corr. structure)
				forval i=1/`M'{
					forval j=1/`M'{
						qui mat `C'[`i',`j']=`V_hom'[`i',`j'] / sqrt(`V_hom'[`i',`i']*`V_hom'[`j',`j'])
					}
				}
				//fill diagonal
				forval i=1/`M'{
					qui mat `V'[`i',`i']=max(`V_hom'[`i',`i'],`V_het'[`i',`i'])
				}
				//fill off-diagonal
				forval i=1/`M'{
					forval j=1/`M'{
						if `i'!=`j' qui mat `V'[`i',`j']=`C'[`i',`j']*sqrt(`V'[`i',`i']*`V'[`j',`j'])
					}
				}
			}


			//this mat for back-compatibility
			mat Jtot = colsof(`b')
			local Jtot = Jtot[1,1]
			local J1 `J'
			local J2 = `Jtot' - `J'

			//recenter within each copy
			mat `R1' = J(`J1',`J1',1/`J1')
			mat `R2' = J(`J2',`J2',1/`J2')
			mata: st_matrix("`R'",blockdiag(st_matrix("`R1'"),st_matrix("`R2'")))
			mat `T' = I(`Jtot') - `R'

			//recenter point ests
			mat `b' = `b'*`T'
			mat colnames `b' = `bnames'

			//recenter VCE
			if "`vce'"!="" mat `V' = `T'*`V'*`T''
			if "`vce'" != ""{
				mat rownames `V' = `: colfullnames e(b)'
				mat colnames `V' = `: colfullnames e(b)'
			}

			mat `b_short' = J(1,`J',.)
			if "`vce'"!=""{
				mat `V_short' = J(`J',`J',.)
			}

			//included enrollment vars
			local Xincluded = e(insts1)
			local i 1
			//use p-score controlled estimate if available, if not use conventional est
			foreach x of varlist `D'{
				local j 1
				if strpos("`Xincluded'","`p_`x''") > 0{
					mat `b_short'[1,`i'] = `b'[1,"`p_`x''"]
					if "`vce'"!=""{
						foreach y of varlist `D'{
							if strpos("`Xincluded'","`p_`y''") > 0  mat `V_short'[`i',`j'] = `V'["`p_`x''","`p_`y''"]
							else mat `V_short'[`i',`j'] = `V'["`p_`x''","`y'"]
							local ++j
						}
					}
				}

				else{
					mat `b_short'[1,`i'] = `b'[1,"`x'"]
					if "`vce'"!=""{
						foreach y of varlist `D'{
							if strpos("`Xincluded'","`p_`y''") > 0 mat `V_short'[`i',`j'] = `V'["`x'","`p_`y''"]
							else mat `V_short'[`i',`j'] = `V'["`x'","`y'"]
							local ++j
						}
					}
				}
				local ++i
			}
		restore
		mat `b' = `b_short'
		if "`vce'"!="" mat `V' = `V_short'
	}

	//matrix for recentering
	mat `T' = I(`J') - J(`J',`J',1/`J')

	//recenter point ests
	mat `b' = `b'*`T'

	//recenter VCE
	if "`vce'"!="" mat `V' = `T'*`V'*`T''

	//store
	tempname enrgroup alphaD varD matD
	egen `enrgroup' = group(`D')

	preserve
		keep `D' `enrgroup'
		qui duplicates drop

		mkmat `D', matrix(`matD')

		mat `alphaD' = `matD' * `b''
		svmat `alphaD', names(`gen')
		rename `gen'1 `gen'

		//yearly VAM's
		mat colnames `b' = `D'
		return mat b = `b'

		//list of imputed VAMs
		if "`pscores'"!=""{
			local imputed: list D - lottD
			return local imputed "`imputed'"
		}

		//full VCE
		if "`vce'"!=""{
			/*
			mat `varD' = vecdiag(`matD' * `V' * `matD'')'
			svmat `varD', names(V_`gen')
			rename V_`gen'1 V_`gen'
			*/
			mat rownames `V' = `D'
			mat colnames `V' = `D'
			return mat V = `V'
		}

		//school counts
		local J: word count `D'
		return local J = `J'
		if "`pscores'"!=""{
			local J_lottery: word count `p_D'
			return local J_lottery = `J_lottery'
		}

		tempfile vam
		qui save "`vam'"
	restore

	qui merge m:1 `enrgroup' using "`vam'", nogen keepusing(`gen')
	di "olsvam: variable `gen' generated"

	end
